// ntp_logic_app.cpp
// 单文件示例：A(请求端)/B(应答端) 使用 LogicalClock 打 t1/t2/t3/t4 并做 NTP 估计
// 仅改动点：1) make_req 接受 tsync::LogicalClock&  2) run_A/run_B 接受 tsync::LogicalClock& 并用它打戳

#include <cstdint>
#include <cstring>
#include <chrono>
#include <algorithm>
#include <vector>
#include <limits>
#include <string>
#include <iostream>
#include <thread>
#include <mutex>

#include "logical_clock.h"  // tsync::LogicalClock
#include "time_source.h"

#if defined(_WIN32)
  #include <winsock2.h>
  #include <ws2tcpip.h>
  #pragma comment(lib, "Ws2_32.lib")
  using socklen_t = int;
#else
  #include <arpa/inet.h>
  #include <sys/socket.h>
  #include <sys/types.h>
  #include <netinet/in.h>
  #include <unistd.h>
#endif

#ifndef NTPSYNC_DEBUG
#define NTPSYNC_DEBUG 1
#endif
#ifndef NTPSYNC_DEBUG_VERBOSE
#define NTPSYNC_DEBUG_VERBOSE 0
#endif
#ifndef NTPSYNC_DEBUG_WIRE
#define NTPSYNC_DEBUG_WIRE 0
#endif

#if NTPSYNC_DEBUG
static std::mutex g_dbg_mtx;
#define DBG(fmt, ...) do { \
  std::lock_guard<std::mutex> _lk(g_dbg_mtx); \
  std::fprintf(stderr, "[DBG] " fmt "\n", ##__VA_ARGS__); \
  std::fflush(stderr); \
} while(0)
#else
#define DBG(...) ((void)0)
#endif

#if NTPSYNC_DEBUG && NTPSYNC_DEBUG_WIRE
#define DBG_WIRE(...) DBG(__VA_ARGS__)
#else
#define DBG_WIRE(...) ((void)0)
#endif

#ifndef __BYTE_ORDER__
#  if defined(__linux__)
#    include <endian.h>
#  endif
#endif
#ifndef __BYTE_ORDER__
#  if defined(_WIN32)
#    define __BYTE_ORDER__ __ORDER_LITTLE_ENDIAN__
#  else
#    define __BYTE_ORDER__ __ORDER_LITTLE_ENDIAN__
#  endif
#endif

static inline uint64_t htonll(uint64_t v){
#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
  return ( (uint64_t)htonl((uint32_t)(v & 0xffffffffULL)) << 32 ) | htonl((uint32_t)(v >> 32));
#else
  return v;
#endif
}
static inline uint64_t ntohll(uint64_t v){
#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
  return ( (uint64_t)ntohl((uint32_t)(v & 0xffffffffULL)) << 32 ) | ntohl((uint32_t)(v >> 32));
#else
  return v;
#endif
}

#pragma pack(push, 1)
struct SyncNtpMsg {
  uint8_t   version{1};
  uint8_t   mode{0};       // 1=REQ 2=RESP
  uint16_t  flags{0};
  uint32_t  k{0};
  uint32_t  phi_s{0};      // 相位(ns, 0..tau)
  int32_t   r_i_ppm{0};
  uint64_t  t1_ns{0};
  uint64_t  t2_ns{0};
  uint64_t  t3_ns{0};

  void hton() {
    DBG_WIRE("hton() v=%u m=%u k=%u (before) t1=%llu t2=%llu t3=%llu",
             version, mode, k,
             (unsigned long long)t1_ns,
             (unsigned long long)t2_ns,
             (unsigned long long)t3_ns);
    uint32_t tmp_u = htonl(static_cast<uint32_t>(r_i_ppm));
    std::memcpy(&r_i_ppm, &tmp_u, sizeof(tmp_u));
    flags = htons(flags);
    k     = htonl(k);
    phi_s = htonl(phi_s);
    t1_ns = htonll(t1_ns);
    t2_ns = htonll(t2_ns);
    t3_ns = htonll(t3_ns);
  }
  void ntoh() {
    uint32_t tmp_u;
    std::memcpy(&tmp_u, &r_i_ppm, sizeof(tmp_u));
    r_i_ppm = static_cast<int32_t>(ntohl(tmp_u));
    flags   = ntohs(flags);
    k       = ntohl(k);
    phi_s   = ntohl(phi_s);
    t1_ns   = ntohll(t1_ns);
    t2_ns   = ntohll(t2_ns);
    t3_ns   = ntohll(t3_ns);
    DBG_WIRE("ntoh() v=%u m=%u k=%u (after) t1=%llu t2=%llu t3=%llu",
             version, mode, k,
             (unsigned long long)t1_ns,
             (unsigned long long)t2_ns,
             (unsigned long long)t3_ns);
  }

  static constexpr size_t size_bytes() { return sizeof(SyncNtpMsg); }

  // === 改动：使用“外部传入”的 tsync::LogicalClock 打 t1/phi ===
  static SyncNtpMsg make_req(uint32_t k, uint32_t tau_ns, int32_t r_i_ppm,
                             tsync::LogicalClock& lc) {
    SyncNtpMsg m;
    m.version = 1;
    m.mode    = 1;
    m.k       = k;
    m.r_i_ppm = r_i_ppm;
    uint64_t now = (uint64_t)lc.now_ns();
    m.t1_ns   = now;
    m.phi_s   = static_cast<uint32_t>(now % tau_ns);
    DBG_WIRE("make_req() k=%u t1=%llu phi=%u ppm=%d",
             k, (unsigned long long)m.t1_ns, m.phi_s, m.r_i_ppm);
    return m;
  }

  static SyncNtpMsg make_resp_from_req(const SyncNtpMsg& req) {
    SyncNtpMsg m;
    m.version = req.version;
    m.mode    = 2;
    m.k       = req.k;
    m.r_i_ppm = 0;
    m.phi_s   = 0;
    m.t1_ns   = req.t1_ns; // 回显对端 T1
    DBG_WIRE("make_resp_from_req() k=%u echo_t1=%llu",
             m.k, (unsigned long long)m.t1_ns);
    return m;
  }
};
#pragma pack(pop)

static_assert(sizeof(SyncNtpMsg) == 40, "SyncNtpMsg wire size mismatch");

struct SyncNtpEstimate {
  bool     ok{false};
  uint32_t k{0};
  int64_t  delay_ns{0};
  int64_t  offset_ns{0};    // (B - A)
  int64_t  min_delay_ns{0};
  uint64_t t1_ns{0};
  uint64_t t2_ns{0};
  uint64_t t3_ns{0};
  uint64_t t4_ns{0};
};

class SyncNtpEstimator {
public:
  explicit SyncNtpEstimator(size_t min_win=32): win_(min_win) {}
  SyncNtpEstimate estimate_from_resp(const SyncNtpMsg& resp_net, uint64_t t4_ns,
                                     uint64_t t1_fallback=0) {
    SyncNtpEstimate est{};
    SyncNtpMsg m = resp_net; m.ntoh();

    uint64_t t1 = (t1_fallback!=0) ? t1_fallback : m.t1_ns;
    est.t1_ns = t1;
    est.t2_ns = m.t2_ns;
    est.t3_ns = m.t3_ns;
    est.t4_ns = t4_ns;

    if (m.mode!=2 || m.t2_ns==0 || m.t3_ns==0 || t1==0 || t4_ns==0) return est;

    const int64_t d1 = (int64_t)m.t2_ns - (int64_t)t1;
    const int64_t d2 = (int64_t)t4_ns  - (int64_t)m.t3_ns;
    const int64_t delay  = (d1 + d2) / 2;
    const int64_t offset = (d1 - d2) / 2; // B - A

    push_min_delay(delay);

    est.ok           = true;
    est.k            = m.k;
    est.delay_ns     = delay;
    est.offset_ns    = offset;
    est.min_delay_ns = min_delay_ns();
    return est;
  }
  int64_t min_delay_ns() const {
    if (buf_.empty()) return std::numeric_limits<int64_t>::max();
    return *std::min_element(buf_.begin(), buf_.end());
  }
  void clear() { buf_.clear(); }
private:
  void push_min_delay(int64_t d) {
    if (d<=0) return;
    buf_.push_back(d);
    if (buf_.size()>win_) buf_.erase(buf_.begin());
  }
  size_t win_;
  std::vector<int64_t> buf_;
};

// ===================== UDP 辅助 =====================
static int udp_bind(const char* ip, uint16_t port) {
  int fd = (int)socket(AF_INET, SOCK_DGRAM, 0);
  if (fd < 0) { perror("socket"); return -1; }
  sockaddr_in addr{}; addr.sin_family=AF_INET; addr.sin_port=htons(port);
  if (inet_pton(AF_INET, ip, &addr.sin_addr) != 1) { perror("inet_pton"); return -1; }
  int on=1; setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, (const char*)&on, sizeof(on));
  if (bind(fd, (sockaddr*)&addr, sizeof(addr)) < 0) { perror("bind"); return -1; }
  return fd;
}
static int udp_peer(const char* ip, uint16_t port, sockaddr_in& dst){
  dst = {}; dst.sin_family=AF_INET; dst.sin_port=htons(port);
  if (inet_pton(AF_INET, ip, &dst.sin_addr) != 1) { perror("inet_pton"); return -1; }
  return 0;
}

// ===================== A 端（请求端） =====================
// === 改动：传入上层统一维护的 tsync::LogicalClock& lc ===
static void run_A(const char* dst_ip, uint16_t port, tsync::LogicalClock& lc) {
#if defined(_WIN32)
  WSADATA wsa; WSAStartup(MAKEWORD(2,2), &wsa);
#endif
  int fd = udp_bind("0.0.0.0", 0);
  if (fd < 0) return;

  sockaddr_in dst{}; udp_peer(dst_ip, port, dst);

#if defined(_WIN32)
  DWORD tv = 50; setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, (const char*)&tv, sizeof(tv));
#else
  timeval tv{}; tv.tv_sec=0; tv.tv_usec=50*1000;
  setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
#endif

  auto cal = TscSource::init(0.2);
  if (!cal.ok) std::fprintf(stderr, "[WARN] TSC init failed, fallback clock.\n");

  SyncNtpEstimator estimator(32);

  uint32_t k = 0;
  const uint32_t tau_ns = 200'000'000; // 200ms 或 500ms；按你需求
  const int32_t  r_i_ppm = 0;

  std::cout << "[A] start, target=" << dst_ip << ":" << port << std::endl;

  int print_cnt = 0;

  for (;;) {
    // t1：用外部 lc
    SyncNtpMsg req = SyncNtpMsg::make_req(++k, tau_ns, r_i_ppm, lc);
    uint64_t t1_local = (uint64_t)lc.now_ns();  // fallback
    SyncNtpMsg req_net = req; req_net.hton();

    if (sendto(fd, (const char*)&req_net, sizeof(req_net), 0, (sockaddr*)&dst, sizeof(dst)) < 0) {
      perror("[A] sendto"); break;
    }
    DBG("[A] sent REQ k=%u t1=%llu -> %s:%u",
        req.k, (unsigned long long)req.t1_ns, dst_ip, (unsigned)port);

    // t4：用外部 lc
    SyncNtpMsg resp_net{}; sockaddr_in src{}; socklen_t sl=sizeof(src);
    ssize_t n = recvfrom(fd, (char*)&resp_net, sizeof(resp_net), 0, (sockaddr*)&src, &sl);
    uint64_t t4_ns = (uint64_t)lc.now_ns();
    if (n < (ssize_t)sizeof(SyncNtpMsg)) {
      std::this_thread::sleep_for(std::chrono::milliseconds(500));
      continue;
    }
    DBG("[A] recv RESP (net order) bytes=%zd t4=%llu", n, (unsigned long long)t4_ns);

    if (print_cnt++ < 5) {
      SyncNtpMsg peek=resp_net; peek.ntoh();
      DBG("[A] resp ntoh k=%u v=%u m=%u t1=%llu t2=%llu t3=%llu t4=%llu",
          peek.k, peek.version, peek.mode,
          (unsigned long long)peek.t1_ns,
          (unsigned long long)peek.t2_ns,
          (unsigned long long)peek.t3_ns,
          (unsigned long long)t4_ns);
    }

    auto est = estimator.estimate_from_resp(resp_net, t4_ns, t1_local);
    if (est.ok) {
      if (std::llabs(est.offset_ns) > 1'000'000'000LL) {
        DBG("[A] drop abnormal offset=%lld ns", (long long)est.offset_ns);
      } else {
        std::cout << "[A] k=" << est.k
                  << " delay="  << est.delay_ns/1e3 << "us"
                  << " offset=" << est.offset_ns    << "ns"
                  << " min_delay=" << est.min_delay_ns/1e3 << "us"
                  << std::endl;
        // 可在此用 lc.add_offset_step(...) 或 lc.apply_mu_* 做校正
      }
    }
    std::this_thread::sleep_for(std::chrono::milliseconds(500));
  }

#if defined(_WIN32)
  closesocket(fd); WSACleanup();
#else
  close(fd);
#endif
}

// ===================== B 端（应答端） =====================
// === 改动：传入上层统一维护的 tsync::LogicalClock& lc ===
static void run_B(const char* bind_ip, uint16_t port, tsync::LogicalClock& lc) {
#if defined(_WIN32)
  WSADATA wsa; WSAStartup(MAKEWORD(2,2), &wsa);
#endif
  int fd = udp_bind(bind_ip, port);
  if (fd < 0) return;

#if !defined(_WIN32)
  timeval tv{}; tv.tv_sec=1; tv.tv_usec=0;
  setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
#endif

  auto cal = TscSource::init(0.2);
  if (!cal.ok) std::fprintf(stderr, "[WARN] TSC init failed, fallback clock.\n");

  std::cout << "[B] listening on " << bind_ip << ":" << port << std::endl;

  int print_cnt = 0;

  for (;;) {
    // 收到 REQ：t2 用外部 lc
    SyncNtpMsg req_net{}; sockaddr_in src{}; socklen_t sl=sizeof(src);
    ssize_t n = recvfrom(fd, (char*)&req_net, sizeof(req_net), 0, (sockaddr*)&src, &sl);
    uint64_t t2_ns_logic = (uint64_t)lc.now_ns();
    if (n < (ssize_t)sizeof(SyncNtpMsg)) continue;

    DBG("[B] recv REQ (net order) bytes=%zd", n);

    SyncNtpMsg req = req_net; req.ntoh();
    if (print_cnt++ < 5) {
      DBG("[B] req ntoh k=%u v=%u m=%u t1=%llu (t2_logic=%llu)",
          req.k, req.version, req.mode,
          (unsigned long long)req.t1_ns,
          (unsigned long long)t2_ns_logic);
    }

    // 构造 RESP：t2=t2_ns_logic，t3=发送前再拍（同一只 lc）
    SyncNtpMsg resp = SyncNtpMsg::make_resp_from_req(req);
    resp.t2_ns = t2_ns_logic;
    resp.t3_ns = (uint64_t)lc.now_ns();
    SyncNtpMsg resp_net = resp; resp_net.hton();

    if (sendto(fd, (const char*)&resp_net, sizeof(resp_net), 0, (sockaddr*)&src, sizeof(src)) < 0) {
      perror("[B] sendto"); break;
    }
    DBG("[B] send RESP k=%u t1=%llu t2=%llu t3=%llu -> back",
        resp.k,
        (unsigned long long)resp.t1_ns,
        (unsigned long long)resp.t2_ns,
        (unsigned long long)resp.t3_ns);
  }

#if defined(_WIN32)
  closesocket(fd); WSACleanup();
#else
  close(fd);
#endif
}

/* 你在上层这样用：
int main(int argc, char** argv) {
  auto cal = TscSource::init(0.2);
  tsync::LogicalClock lc(0); // 统一维护的一只时钟
  if (argv[1][0]=='A') run_A(argv[2], (uint16_t)std::stoi(argv[3]), lc);
  else                 run_B(argv[2], (uint16_t)std::stoi(argv[3]), lc);
}
*/
